/*
iovpn
*/
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"sync"

	"github.com/songgao/water"
	"github.com/songgao/water/waterutil"
)

type config struct {
	Proto      string
	Devtype    string
	Devnetwork string
	Localaddr  string
	Peeraddrs  string
}

type Table struct {
	Ipmap sync.Map
	Peers []string
}

func (t *Table) AddElms(peer ...string) {
	ps := make(map[string]bool, 0)
	tmp := make([]string, 0)
	for _, v := range t.Peers {
		ps[v] = true
	}
	for _, v := range peer {
		ps[v] = true
	}
	for k := range ps {
		tmp = append(tmp, k)
	}
	t.Peers = tmp
}

func (t *Table) DeleteElms(peer ...string) {
	ps := make(map[string]bool, 0)
	tmp := make([]string, 0)
	for _, v := range t.Peers {
		ps[v] = true
	}
	for _, v := range peer {
		delete(ps, v)
	}
	t.Peers = tmp
}

func main() {
	var Conf = &config{}
	flag.StringVar(&Conf.Proto, "p", "udp", "proto of transfor")
	flag.StringVar(&Conf.Devtype, "if", "tun", "type of tuntap device[tap|tun] ")
	flag.StringVar(&Conf.Devnetwork, "net", "10.0.0.1/32", "ip address of tun/tap")
	flag.StringVar(&Conf.Peeraddrs, "peer", "", "peer addresses used [,] split")
	flag.StringVar(&Conf.Localaddr, "l", "0.0.0.0:8001", "local  address")
	flag.Parse()
	switch Conf.Proto {
	case "http":
		TransHttp(Conf, createtuntap(Conf.Devtype, Conf.Devnetwork))
	default:
		TransUdp(Conf, createtuntap(Conf.Devtype, Conf.Devnetwork))
	}
}

func createtuntap(devtype string, devnetwork string) water.Interface {
	fmt.Println(SetColor("1", "32", "[init "+os.Args[0]+"]"))
	defer fmt.Println(SetColor("1", "32", "[start "+os.Args[0]+"]"))
	// 选择tuntap类型，及设置配置
	config := water.Config{}
	switch devtype {
	case "tap":
		config.DeviceType = water.TAP
	default:
		config.DeviceType = water.TUN
	}
	// 根据配置生成接口
	ifce, err := water.New(config)
	if err != nil {
		log.Fatal(err)
	}
	err = setupNetDevice("addr", "add", devnetwork, "dev", ifce.Name())
	if err != nil {
		fmt.Println("[", ifce.Name(), "]:设置网络地址失败：", err)
	} else {
		fmt.Println("[", ifce.Name(), "]:设置地址：", devnetwork)
	}
	err = setupNetDevice("link", "set", ifce.Name(), "mtu", "1400", "up")
	if err != nil {
		fmt.Println("[", ifce.Name(), "]:启动失败：", err)
	} else {
		fmt.Println("[", ifce.Name(), "]:启动成功")
	}
	return *ifce
}

func setupNetDevice(cmd ...string) error {
	cmdexec := exec.Command("ip", cmd...)
	cmdexec.Stderr = os.Stderr
	cmdexec.Stdout = os.Stdout
	err := cmdexec.Run()
	if err != nil {
		return err
	}
	return nil
}

/*udp*/
func TransUdp(conf *config, ifio water.Interface) {
	defer ifio.Close()
	var tab = &Table{Peers: make([]string, 0)}
	//1.解析本地地址
	locaddr, err := net.ResolveUDPAddr("udp", conf.Localaddr)
	if err != nil {
		log.Fatalln(SetColor("2", "31", "parse udp address error:"), err)
		return
	}
	//2.解析公网节点地址
	tab.Peers = strings.Split(conf.Peeraddrs, ",")
	if tab.Peers[0] != "" {
		for _, peer := range tab.Peers {
			paddr, err := net.ResolveUDPAddr("udp", peer)
			if err != nil {
				log.Fatalln(SetColor("2", "31", "parse udp address error:"), err)
				return
			}
			fmt.Println(SetColor("2", "32", "[Add peer]: <== "), paddr.String())
		}
	}
	//3.监听udp
	listener, err := net.ListenUDP("udp", locaddr)
	if err != nil {
		log.Fatalln(err)
		return
	}
	// udp --> device
	go func() {
		buf := make([]byte, 1500)
		for {
			n, clientaddr, err := listener.ReadFromUDP(buf)
			if err != nil {
				fmt.Println(SetColor("3", "31", "resv udp error:"), err)
				continue
			}
			if ifio.IsTAP() {
				dstMAC := waterutil.MACDestination(buf[:n])
				switch waterutil.MACEthertype(buf[:n]) {
				case waterutil.ARP:
					ARPsIP := ARPsrcIP(buf[:n])
					if waterutil.IsBroadcast(dstMAC) {
						v, ok := tab.Ipmap.LoadOrStore(ARPsIP.String(), clientaddr.String())
						if !ok {
							fmt.Println(SetColor("2", "32", "[Ipmap]:add "), ARPsIP.String(), "==>", clientaddr.String())
						} else {
							if v.(string) != clientaddr.String() {
								tab.Ipmap.Store(ARPsIP.String(), clientaddr.String())
							}
						}
						ifio.Write(buf[:n])
					} else {
						v, ok := tab.Ipmap.LoadOrStore(ARPsIP.String(), clientaddr.String())
						if !ok {
							fmt.Println(SetColor("2", "32", "[Ipmap]:add "), ARPsIP.String(), "==>", clientaddr.String())
						} else {
							if v.(string) != clientaddr.String() {
								tab.Ipmap.Store(ARPsIP.String(), clientaddr.String())
							}
						}
						ifio.Write(buf[:n])
					}
				case waterutil.IPv4:
					ifio.Write(buf[:n])
				}
			} else {
				if waterutil.IsIPv4(buf[:n]) {
					srcIP := waterutil.IPv4Source(buf[:n])
					v, ok := tab.Ipmap.LoadOrStore(srcIP.String(), clientaddr.String())
					if !ok {
						fmt.Println(SetColor("2", "32", "[Ipmap]:add "), srcIP.String(), "==>", clientaddr.String())
					} else {
						if v.(string) != clientaddr.String() {
							tab.Ipmap.Store(srcIP.String(), clientaddr.String())
						}
					}
					ifio.Write(buf[:n])
				}
			}
		}
	}()
	// device --> udp
	func() {
		// // udp心跳
		// var heartbeat = [3]byte{'u', 'd', 'p'}
		// go func() {
		// 	for {
		// 		for _, peer := range tab.Peers {
		// 			paddr, err := net.ResolveUDPAddr("udp", peer)
		// 			if err != nil {
		// 				fmt.Println(SetColor("3", "31", "parse udp address error:"), err)
		// 				return
		// 			}
		// 			listener.WriteToUDP(heartbeat[:], paddr)
		// 		}
		// 		time.Sleep(3 * time.Second)
		// 	}
		// }()
		buf := make([]byte, 1500)
		for {
			n, err := ifio.Read(buf)
			if err != nil {
				log.Fatalln(err)
				continue
			}
			if ifio.IsTAP() {
				//	srcMAC := waterutil.MACSource(buf[:n])
				dstMAC := waterutil.MACDestination(buf[:n])
				switch waterutil.MACEthertype(buf[:n]) {
				case waterutil.ARP:
					//ARPsIP := ARPsrcIP(buf[:n])
					ARPdIP := ARPdstIP(buf[:n])
					if waterutil.IsBroadcast(dstMAC) {
						for _, peer := range tab.Peers {
							paddr, err := net.ResolveUDPAddr("udp", peer)
							if err != nil {
								log.Fatalln(err)
								continue
							}
							_, err = listener.WriteToUDP(buf[:n], paddr)
							if err != nil {
								fmt.Println(SetColor("3", "31", "send upd error:"), err)
								continue
							}
						}
					} else {
						//fmt.Println("[ARP]resp src:", srcMAC.String(), "dst:", dstMAC.String(), "resp ip:", ARPsIP.String())
						raddr, ok := tab.Ipmap.Load(ARPdIP.String())
						if ok {
							peer, ok := raddr.(string)
							if ok {
								paddr, err := net.ResolveUDPAddr("udp", peer)
								if err != nil {
									fmt.Println(SetColor("3", "31", "parse udp address error:"), err)
									continue
								}
								_, err = listener.WriteToUDP(buf[:n], paddr)
								if err != nil {
									fmt.Println(SetColor("3", "31", "send udp error:"), err)
									continue
								}
							} else {
								fmt.Println(SetColor("3", "31", "parse error"), raddr)
								continue
							}
						} else {
							fmt.Println(SetColor("3", "31", "found Ipmap error"))
							continue
						}
					}
				case waterutil.IPv4:
					//fmt.Println("[IPv4]src:", srcMAC.String(), "-->", "dst:", dstMAC.String())
					//fmt.Println("[IPv4]dest ip:", waterutil.IPv4Destination(waterutil.MACPayload(buf[:n])).String())
					v, ok := tab.Ipmap.Load(waterutil.IPv4Destination(waterutil.MACPayload(buf[:n])).String())
					if !ok {
						fmt.Println(SetColor("3", "31", "not found ip in map:"), waterutil.IPv4Destination(waterutil.MACPayload(buf[:n])).String())
						continue
					} else {
						raddr, ok := v.(string)
						if ok {
							remoteaddr, err := net.ResolveUDPAddr("udp", raddr)
							if err != nil {
								fmt.Println(SetColor("3", "31", "parse udp address error:"), err)
								continue
							}
							_, err = listener.WriteToUDP(buf[:n], remoteaddr)
							if err != nil {
								fmt.Println(SetColor("3", "31", "send upd error:"), err)
								continue
							}
						}
					}
				}
			} else {
				dstIP := waterutil.IPv4Destination(buf[:n])
				v, ok := tab.Ipmap.Load(dstIP.String())
				if ok {
					remote, err := net.ResolveUDPAddr("udp", v.(string))
					if err != nil {
						tab.Ipmap.Delete(dstIP.String())
						fmt.Println(SetColor("3", "31", "[Delete IPmap]: "), "==> ", v.(string), err)
						continue
					}
					_, err = listener.WriteToUDP(buf[:n], remote)
					if err != nil {
						fmt.Println(err)
						continue
					}
				} else {
					if len(tab.Peers) == 1 && tab.Peers[0] == "" {
						tab.Ipmap.Range(func(key, value any) bool {
							remote, err := net.ResolveUDPAddr("udp", value.(string))
							if err != nil {
								tab.Ipmap.Delete(key)
								fmt.Println(SetColor("3", "31", "[Delete IPmap]: "), "==> ", value.(string), err)
							}
							_, err = listener.WriteToUDP(buf[:n], remote)
							if err != nil {
								tab.Ipmap.Delete(key)
								fmt.Println(SetColor("2", "31", "[Delete IPmap] send udp error: "), "==> ", value.(string), err)
								return true
							}
							return true
						})
					} else {
						for _, peer := range tab.Peers {
							remote, err := net.ResolveUDPAddr("udp", peer)
							if err != nil {
								fmt.Println(err)
								continue
							}
							_, err = listener.WriteToUDP(buf[:n], remote)
							if err != nil {
								fmt.Println(err)
								continue
							}
						}
					}
				}
			}
		}
	}()
}

// color tools
func SetColor(i, c, in string) string {
	return fmt.Sprint("\033[", i, ";", c, "m", in, "\033[0m")
}

// ARP工具
func ARPsrcIP(frame []byte) net.IP {
	return net.IPv4(frame[28], frame[29], frame[30], frame[31])
}

func ARPdstIP(frame []byte) net.IP {
	return net.IPv4(frame[38], frame[39], frame[40], frame[41])
}

//[255 255 255 255 255 255 34 210 66 114 118 138 8 6 0 1 8 0 6 4 0 1 34 210 66 114 118 138 10 0 0 4 0 0 0 0 0 0 10 0 0 9]
//                  dstmac                srcmac12131415161718192021                srcmac 28                   38

func post(url, addr string, data []byte) (*http.Response, error) {
	req, _ := http.NewRequest("POST", url, bytes.NewReader(data))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("remoter", addr)
	return http.DefaultClient.Do(req)
}

/*http*/
func TransHttp(conf *config, ifio water.Interface) {
	defer ifio.Close()
	var tab = &Table{Peers: make([]string, 0)}
	// 初始化tab
	tab.Peers = strings.Split(conf.Peeraddrs, ",")
	if len(tab.Peers) > 0 {
		for _, peer := range tab.Peers {
			paddr, err := net.ResolveTCPAddr("tcp", peer)
			if err != nil {
				log.Fatalln(err, "uuu")
				return
			}
			fmt.Println("[Add peer]:==> ", paddr.String())
		}
	}
	//net --> dev
	go func() {
		buf := make([]byte, 1500)
		http.HandleFunc("/net", func(w http.ResponseWriter, r *http.Request) {
			if remoter := r.Header.Get("remoter"); remoter != "" && r.Method == "POST" {
				n, err := bufio.NewReader(r.Body).Read(buf)
				if err != nil && err != io.EOF {
					fmt.Println("jjj", err)
				}
				if ifio.IsTAP() {
					dstMAC := waterutil.MACDestination(buf[:n])
					switch waterutil.MACEthertype(buf) {
					case waterutil.ARP:
						ARPsIP := ARPsrcIP(buf[:n])
						switch waterutil.IsBroadcast(dstMAC) {
						case true:
							_, ok := tab.Ipmap.LoadOrStore(ARPsIP.String(), remoter)
							if !ok {
								fmt.Println("[Ipmap]:add", ARPsIP.String(), "==>", remoter)
								tab.AddElms(remoter)
							}
							ifio.Write(buf[:n])
							w.WriteHeader(http.StatusOK)
							w.Write([]byte(conf.Localaddr))
						default:
							_, ok := tab.Ipmap.LoadOrStore(ARPsIP.String(), remoter)
							if !ok {
								fmt.Println("[Ipmap]:add", ARPsIP.String(), "==>", remoter)
								tab.AddElms(remoter)
							}
							ifio.Write(buf[:n])
							w.WriteHeader(http.StatusOK)
							w.Write([]byte(conf.Localaddr))
						}
					case waterutil.IPv4:
						ifio.Write(buf[:n])
						w.WriteHeader(http.StatusOK)
						w.Write([]byte(conf.Localaddr))
					}
				} else {
					if waterutil.IsIPv4(buf[:n]) {
						srcIP := waterutil.IPv4Source(buf[:n])
						_, ok := tab.Ipmap.LoadOrStore(srcIP.String(), remoter)
						if !ok {
							fmt.Println("[Ipmap]:add", srcIP.String(), "==> ", remoter)
							tab.AddElms(remoter)
						}
						ifio.Write(buf[:n])
						w.WriteHeader(http.StatusOK)
						w.Write([]byte(conf.Localaddr))
					}
				}
			} else {
				w.WriteHeader(http.StatusOK)
				w.Write([]byte(conf.Localaddr))
			}
		})
		log.Fatalln(http.ListenAndServe(conf.Localaddr, nil))
	}()
	//dev --> net
	func() {
		buf := make([]byte, 1500)
		for {
			n, err := ifio.Read(buf)
			if err != nil {
				fmt.Println(err)
				continue
			}
			if ifio.IsTAP() {
				dstMAC := waterutil.MACDestination(buf[:n])
				switch waterutil.MACEthertype(buf[:n]) {
				case waterutil.ARP:
					ARPdIP := ARPdstIP(buf[:n])
					if waterutil.IsBroadcast(dstMAC) {
						for _, peer := range tab.Peers {
							post(fmt.Sprint("http://", peer, "/net"), conf.Localaddr, buf[:n])
						}
					} else {
						raddr, ok := tab.Ipmap.Load(ARPdIP.String())
						if ok {
							peer, ok := raddr.(string)
							if ok {
								post(fmt.Sprint("http://", peer, "/net"), conf.Localaddr, buf[:n])
							} else {
								fmt.Println("解析失败", raddr)
								continue
							}
						} else {
							fmt.Println("查router失败")
							continue
						}
					}
				case waterutil.IPv4:
					v, ok := tab.Ipmap.Load(waterutil.IPv4Destination(waterutil.MACPayload(buf[:n])).String())
					if !ok {
						fmt.Println("[Ipmap]net found ip in map:", waterutil.IPv4Destination(waterutil.MACPayload(buf[:n])).String())
						continue
					} else {
						peer, ok := v.(string)
						if ok {
							post(fmt.Sprint("http://", peer, "/net"), conf.Localaddr, buf[:n])
						}
					}
				}
			} else {
				dstIP := waterutil.IPv4Destination(buf[:n])
				v, ok := tab.Ipmap.Load(dstIP.String())
				if ok {
					post(fmt.Sprint("http://", v.(string), "/net"), conf.Localaddr, buf[:n])
				} else {
					for _, peer := range tab.Peers {
						post(fmt.Sprint("http://", peer, "/net"), conf.Localaddr, buf[:n])
					}
				}
			}
		}
	}()
}
